"""功耗采样后端模块"""

from __future__ import annotations

from typing import Dict, Type, List, Optional

from .base import PowerBackend

# 后端注册表：将后端名称映射到后端类
_BACKEND_REGISTRY: Dict[str, Type[PowerBackend]] = {}


def register_backend(name: str, backend_class: Type[PowerBackend]):
    """
    注册功耗后端
    
    Args:
        name: 后端名称
        backend_class: 后端类（必须继承 PowerBackend）
    """
    if not issubclass(backend_class, PowerBackend):
        raise TypeError(f"{backend_class} must inherit from PowerBackend")
    _BACKEND_REGISTRY[name.lower()] = backend_class


def get_power_backend(name: str) -> PowerBackend:
    """
    根据后端名称创建后端实例
    
    Args:
        name: 后端名称（如 "nvml", "rocm", "rapl"）
    
    Returns:
        PowerBackend 实例
    
    Raises:
        ValueError: 后端名称未注册
    """
    name_lower = name.lower()
    
    if name_lower not in _BACKEND_REGISTRY:
        available = ", ".join(sorted(_BACKEND_REGISTRY.keys()))
        raise ValueError(
            f"Unsupported power backend: '{name}'\n"
            f"Available backends: {available}\n"
            f"Use 'list' to see available backends"
        )
    
    backend_class = _BACKEND_REGISTRY[name_lower]
    return backend_class()


def list_power_backends(include_unavailable: bool = True) -> List[str]:
    """
    列出所有已注册的功耗后端
    
    Args:
        include_unavailable: 是否包含不可用的后端
    
    Returns:
        后端名称列表
    """
    backends = sorted(_BACKEND_REGISTRY.keys())
    
    if include_unavailable:
        return backends
    
    # 只返回可用的后端
    available = []
    for name in backends:
        backend_class = _BACKEND_REGISTRY[name]
        if backend_class.is_available():
            available.append(name)
    
    return available


# 导入所有后端以触发注册
# 注意：这些导入只是导入类定义，不会触发外部依赖检查
# 外部依赖（如 pynvml）是在方法内部动态导入的，只有在调用
# is_available()、init() 或 sample() 时才会检查依赖
from .nvml import NVMLBackend
from .rocm import ROCmBackend
from .rapl import RAPLBackend
from .ascend import AscendBackend
from .external_meter import ExternalMeterBackend

# 注册所有后端
register_backend("nvml", NVMLBackend)
register_backend("rocm", ROCmBackend)
register_backend("rapl", RAPLBackend)
register_backend("ascend", AscendBackend)
register_backend("external_meter", ExternalMeterBackend)

__all__ = [
    "PowerBackend",
    "get_power_backend",
    "list_power_backends",
    "register_backend",
    "NVMLBackend",
    "ROCmBackend",
    "RAPLBackend",
    "AscendBackend",
    "ExternalMeterBackend",
]

